import logging
import re
from typing import List
from llama_index.core.schema import Document, BaseNode, TextNode, NodeRelationship
from .markdown_node_parser_config import MarkdownNodeParserConfig
from ..node_parser_interface import NodeParserInterface


class MarkdownNodeParser(NodeParserInterface):
    """
    一个定制化的、实现了 NodeParserInterface 的 Markdown 解析器。
    它能够识别标准 Markdown 标题 (`#`, `##` 等) 和中文特有的章节标题
    (如 `一、`, `（一）`, `1.` 等)，并据此对文档进行切分。
    """

    def __init__(self, config: MarkdownNodeParserConfig):
        self._include_prev_next_rel = config.include_prev_next_rel
        # 正则表达式，用于捕获所有类型的标题作为分隔符
        self._split_pattern = re.compile(
            r'(^#+\s.*|^\s*[一二三四五六七八九十]+、.*|^\s*（[一二三四五六七八九十]+）.*|^\s*[0-9]+\..*)',
            re.MULTILINE
        )

    def parse_nodes(self, documents: List[Document]) -> List[BaseNode]:
        """
        使用定制的分割逻辑来执行节点解析。
        """
        all_nodes: List[BaseNode] = []
        for doc in documents:
            # [诊断日志] 打印传入的 Document 的元数据
            logging.info(f"[METADATA_TRACE] MarkdownParser received doc with metadata: {doc.metadata}")

            text = doc.text
            # 使用 finditer 找到所有匹配的标题及其位置
            matches = list(self._split_pattern.finditer(text))
            if not matches:
                # 如果没有找到任何标题，则将整个文档视为一个节点
                node = TextNode(text=text, metadata=doc.metadata)
                # [最终修复] 必须为所有手动创建的节点添加 SOURCE 关系
                node.relationships[NodeRelationship.SOURCE] = doc.as_related_node_info()
                all_nodes.append(node)
                continue

            # 根据标题位置切分文本
            chunks = []
            last_end = 0
            for match in matches:
                chunk_start = last_end
                chunk_end = match.start()
                chunks.append(text[chunk_start:chunk_end])
                last_end = match.end()
            # 添加最后一个标题之后的内容
            chunks.append(text[last_end:])

            # 将切分后的文本块与标题重新组合成节点
            # 第一个块是没有前置标题的内容
            doc_nodes = []
            if chunks[0].strip():
                node = TextNode(text=chunks[0].strip(), metadata=doc.metadata)
                # [最终修复] 必须为所有手动创建的节点添加 SOURCE 关系
                node.relationships[NodeRelationship.SOURCE] = doc.as_related_node_info()
                doc_nodes.append(node)

            # 从第二个块开始，每个块都由其前面的标题引导
            for i, match in enumerate(matches):
                header = match.group(0).strip()
                content = chunks[i + 1].strip()

                # [最终修复] 必须从源文档继承元数据，而不是创建一个新字典
                # 这确保了 document_id 等关键信息不会丢失
                metadata = doc.metadata.copy()
                metadata['header'] = header

                # [诊断日志] 打印即将创建的 TextNode 的元数据
                logging.info(f"[METADATA_TRACE] MarkdownParser creating node with metadata: {metadata}")

                # [最终修复] 确保节点文本的正确性：
                # - 如果有内容，则文本是 "标题\n\n内容"
                # - 如果没有内容（只有标题），则文本就是标题本身
                node_text = f"{header}\n\n{content}" if content else header
                
                node = TextNode(text=node_text, metadata=metadata)
                # [核心修复] 必须手动添加节点到源文档的 SOURCE 关系
                # 这是 llama-index 内部正确传递元数据的关键
                node.relationships[NodeRelationship.SOURCE] = doc.as_related_node_info()
                doc_nodes.append(node)

            # 建立节点间的关系
            if self._include_prev_next_rel:
                for i, node in enumerate(doc_nodes):
                    if i > 0:
                        node.relationships[NodeRelationship.PREVIOUS] = doc_nodes[i - 1].as_related_node_info()
                    if i < len(doc_nodes) - 1:
                        node.relationships[NodeRelationship.NEXT] = doc_nodes[i + 1].as_related_node_info()

            all_nodes.extend(doc_nodes)
            
        return all_nodes 